from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 诊断相关基础数据 ==========
CANCER_DIAGNOSES = [
    ("C34.90", "肺恶性肿瘤", "ICD-10", "主要诊断", "T2N1M0"),
    ("C50.90", "乳腺恶性肿瘤", "ICD-10", "主要诊断", "T1N0M0"),
    ("C16.90", "胃恶性肿瘤", "ICD-10", "主要诊断", "T3N2M0"),
    ("C18.90", "结肠恶性肿瘤", "ICD-10", "主要诊断", "T4N1M1"),
    ("C22.00", "肝细胞癌", "ICD-10", "主要诊断", "T2N0M0"),
    ("D63.1*", "肿瘤性贫血", "ICD-10", "其他诊断", None),
    ("J18.9", "肺炎", "ICD-10", "医院感染", None),
    ("E11.9", "2型糖尿病", "ICD-10", "并发症", None)
]

DIAG_TYPES = ["入院诊断", "出院诊断", "病理诊断", "并发症诊断", "医院感染诊断"]
OUTCOMES = ["治愈", "好转", "未愈", "死亡", "其他"]
ADMISSION_CONDITIONS = ["危", "急", "一般"]

HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 10 if date_obj.day % 2 == 0 else 6


# ========== 核心生成函数 ==========
def generate_diagnosis_records():
    sql_statements = []
    record_count = 0

    patient_id_counter = 2001
    visit_sn_counter = 2001
    diag_id_counter = 50001
    record_sn_counter = 10001

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1
            patient_id = f"PT{patient_id_counter}"
            patient_id_counter += 1

            visit_sn = f"VIS{visit_sn_counter}"
            visit_sn_counter += 1

            # 随机选择医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 基础信息
            medical_record_no = f"MR{random.randint(100000, 999999)}"
            record_sn = f"REC{record_sn_counter}"
            record_sn_counter += 1

            # 时间信息
            admit_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d") - timedelta(days=30),
                end_date=datetime.strptime(day, "%Y-%m-%d")
            ).strftime('%Y-%m-%d %H:%M:%S')

            hospital_days = random.randint(3, 30)
            discharge_datetime = (datetime.strptime(admit_datetime, '%Y-%m-%d %H:%M:%S') +
                                  timedelta(days=hospital_days)).strftime('%Y-%m-%d %H:%M:%S')

            record_datetime = discharge_datetime
            yy_collection_datetime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            yy_etl_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # 生成诊断记录（每个患者1-4个诊断）
            diag_count = random.randint(1, 4)
            for diag_seq in range(1, diag_count + 1):
                diag_id = f"DIAG{diag_id_counter}"
                diag_id_counter += 1

                # 第一个诊断为主要诊断
                is_main = "是" if diag_seq == 1 else "否"

                # 随机选择诊断（主要诊断必须为肿瘤诊断）
                if diag_seq == 1:
                    diagnosis_code, diagnosis_name, diag_system, diag_type, tnm_staging = random.choice(
                        [d for d in CANCER_DIAGNOSES if d[3] == "主要诊断"]
                    )
                else:
                    diagnosis_code, diagnosis_name, diag_system, diag_type, tnm_staging = random.choice(
                        CANCER_DIAGNOSES)

                # 其他字段
                admission_condition = random.choice(ADMISSION_CONDITIONS)
                out_come = random.choice(OUTCOMES) if diag_seq == 1 else None
                admit_number = str(random.randint(1, 5))

                # 生成yy_record_md5 (模拟MD5)
                yy_record_md5 = f"md5_{random.getrandbits(128):032x}"
                yy_batch_time = day
                yy_record_batch_id = f"BATCH{day}_{i}_{diag_seq}"

                data = {
                    'admission_condition': admission_condition,
                    'admit_datetime': admit_datetime,
                    'admit_number': admit_number,
                    'diag_id': diag_id,
                    'diag_seq': str(diag_seq),
                    'diag_type': diag_type,
                    'diagnosis_code': diagnosis_code,
                    'diagnosis_name': diagnosis_name,
                    'discharge_datetime': discharge_datetime,
                    'from_table': 'EMR_DIAGNOSIS',
                    'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                    'hospital_code': hospital_code,
                    'hospital_name': hospital_name,
                    'is_main': is_main,
                    'medical_record_no': medical_record_no,
                    'out_come': out_come,
                    'patient_id': patient_id,
                    'patient_id_old': f"OLD_{patient_id}",
                    'record_datetime': record_datetime,
                    'record_sn': record_sn,
                    'record_source': 'HIS',
                    'record_status': 1,
                    'record_type': '住院病案',
                    'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'visit_sn': visit_sn,
                    'yy_collection_datetime': yy_collection_datetime,
                    'yy_record_md5': yy_record_md5,
                    'yy_upload_status': 0,
                    'yy_etl_time': yy_etl_time,
                    'yy_upload_time': None,
                    'yy_batch_time': yy_batch_time,
                    'yy_record_batch_id': yy_record_batch_id,
                    'yy_backfill_time': None,
                    'yy_backfill_status': None,
                    'branch_code': branch_code,
                    'branch_name': branch_name,
                    'date_for_partition': admit_datetime
                }

                sql = _generate_sql('b04_1_1', data)
                sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_diagnosis_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b04_1_1')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b04_1_1_diagnosis_records.sql")